import pandas as pd
from .paramsPipe import ParamsPipe
from .dataLoader import DataLoader
import numpy as np


class PortTrack:
    def __init__(self, data: DataLoader, params: ParamsPipe, init_cash):
        self.port_value = init_cash  # 组合总价值
        self.cash = init_cash  # 现金头寸
        # 追踪当前持仓记录, columns=[code, holding_num, holding_price]
        self.holdings = pd.DataFrame(columns=['code', 'holding_num', 'holding_price'])
        self.trading_records = pd.DataFrame()  # 记录交易情况

        self.data = data  # 数据管线
        self.params = params  # 参数传递管线

        # 持仓字典
        self.holdings_dict = {}

    def update_cash(self, cash):
        self.cash = cash

    def add_holding(self, code, size, price):
        # 新增一行数据到holdings DataFrame
        new_holding = pd.DataFrame({
            'code': [code],
            'holding_num': [size],
            'holding_price': [price]  # 以成交价新增持仓信息
        })

        if self.holdings.empty:
            self.holdings = new_holding
        else:
            self.holdings = pd.concat([self.holdings, new_holding], ignore_index=True)

        # 更新字典缓存(合并相同code的持仓)
        if code in self.holdings_dict:
            # 如果已存在该code，合并持仓数量和价格(按持仓数量加权平均)
            existing = self.holdings_dict[code]
            total_num = existing['holding_num'] + size
            if total_num != 0:
                avg_price = (existing['holding_num'] * existing['holding_price'] + size * price) / total_num
                self.holdings_dict[code] = {'holding_num': total_num, 'holding_price': avg_price}
            else:
                # 持仓为0，删除该记录
                del self.holdings_dict[code]
        else:
            self.holdings_dict[code] = {'holding_num': size, 'holding_price': price}

    def sell_holding(self, code, size, price):
        # 卖出后修改holdings, 并删除持有数据数量为0的股票
        mask = self.holdings['code'] == code
        if mask.any():
            # 处理所有匹配的记录
            holding_rows = self.holdings.loc[mask].copy()
            total_holding = holding_rows['holding_num'].sum()

            # 按照先进先出原则处理卖出
            remaining_size = size
            for idx in holding_rows.index:
                if remaining_size <= 0:
                    break
                current_holding = self.holdings.loc[idx, 'holding_num']
                if remaining_size >= current_holding:
                    # 全部卖出该条记录
                    self.holdings.loc[idx, 'holding_num'] = 0
                    remaining_size -= current_holding
                else:
                    # 部分卖出该条记录
                    self.holdings.loc[idx, 'holding_num'] = int(current_holding - remaining_size)
                    self.holdings.loc[idx, 'holding_price'] = price
                    remaining_size = 0

            # 删除持仓为0的记录
            self.holdings = self.holdings[self.holdings['holding_num'] != 0].reset_index(drop=True)

            # 更新字典缓存
            new_total_holding = total_holding - size
            if new_total_holding == 0:
                if code in self.holdings_dict:
                    del self.holdings_dict[code]
            else:
                # 重新计算平均持仓价格
                remaining_value = 0
                remaining_num = 0
                for idx in self.holdings[self.holdings['code'] == code].index:
                    remaining_value += self.holdings.loc[idx, 'holding_num'] * self.holdings.loc[idx, 'holding_price']
                    remaining_num += self.holdings.loc[idx, 'holding_num']

                if remaining_num > 0:
                    avg_price = remaining_value / remaining_num
                    self.holdings_dict[code] = {'holding_num': remaining_num, 'holding_price': avg_price}
                elif code in self.holdings_dict:
                    del self.holdings_dict[code]

    def trading_track(self, date, code, direction, size, price):
        # 更新交易记录，dataframe, columns=[trading, code, direction, size, execute_price]
        new_trading_record = pd.DataFrame({
            'trading': [date],
            'code': [code],
            'direction': [direction],
            'size': [size],
            'price': [price],
        })
        self.trading_records = pd.concat([self.trading_records, new_trading_record], ignore_index=True)

    def _get_today_close(self, code):
        # 获取当前日期，code的收盘价
        _x = self.data.get_close(self.params.today, code)
        if _x is None or np.isnan(_x):
            # 无交易数据，用holding_price代替（多条记录则使用平均成本替代）
            if code in self.holdings_dict:
                _x = self.holdings_dict[code]['holding_price']
            else:
                # fallback到原方法
                _x = self.holdings.loc[self.holdings['code'] == code]['holding_price'].values[0]
        return _x

    def update_portfolio_value(self):
        # 每数据帧结束更新组合总价值
        if self.holdings.empty:
            self.port_value = self.cash
        else:
            # 批量更新价格以提高性能
            self.holdings['holding_price'] = self.holdings['code'].apply(
                lambda code: self._get_today_close(code)
            )
            self.port_value = self.cash + (self.holdings['holding_price'] * self.holdings['holding_num']).sum()

            # 更新字典缓存
            for _, row in self.holdings.iterrows():
                self.holdings_dict[row['code']] = {
                    'holding_num': row['holding_num'],
                    'holding_price': row['holding_price']
                }

    def contribution(self):
        # 计算整个策略期间个股对组合收益率的贡献度
        _total_return = self.port_value - self.params.init_cash
        _raw = self.trading_records.copy()

        result = _raw.groupby('code').agg(
            holding_return=('direction', lambda x: (x * _raw.loc[x.index, 'size'] * _raw.loc[x.index, 'price']).sum())
        ).reset_index()
        result['contribution'] = result['holding_return'] / _total_return
        result.sort_values(by='contribution', ascending=False, inplace=True)
        return result
